import asyncio
import json
import time
from datetime import datetime
import aiohttp
from src.config.config_loader import ConfigLoader

class Benchmark:
    def __init__(self, config_loader: ConfigLoader):
        self.config_loader = config_loader
        # 设置默认超时时间（秒）
        self.timeout = aiohttp.ClientTimeout(total=600)  # 10分钟总超时
        self.connect_timeout = 600  # 30秒连接超时
        self.read_timeout = 600    # 270秒读取超时

    async def send_request(self, session, question, request_id, semaphore):
        """Send a single request to the vLLM server and track performance metrics"""
        headers = self.config_loader.get_headers()
        request_body = {
            "model": self.config_loader.model_name,
            "messages": [{"role": "user", "content": question}],
            "stream": self.config_loader.stream_enabled
        }
        
        # 打印请求信息
        # print(f"\n{'=' * 50}")
        # print(f"请求 {request_id} 详细信息:")
        # print(f"{'=' * 50}")
        # print(f"请求URL: {self.config_loader.server_url}")
        # print(f"请求头:")
        # for key, value in headers.items():
        #     print(f"  {key}: {value}")
        # print(f"请求体:")
        # print(json.dumps(request_body, ensure_ascii=False, indent=2))
        # print(f"{'=' * 50}\n")
        
        start_time = time.time()
        received_tokens = 0
        print(f"[{datetime.now().strftime('%H:%M:%S')}] 请求 {request_id} 已开始: {question[:50]}...")
        
        try:
            # 测试服务器连接
            # try:
            #     async with session.get(self.config_loader.server_url) as test_response:
            #         print(f"[{datetime.now().strftime('%H:%M:%S')}] 服务器连接测试状态码: {test_response.status}")
            # except Exception as e:
            #     print(f"[{datetime.now().strftime('%H:%M:%S')}] 服务器连接测试失败: {str(e)}")
            
            async with semaphore:
                try:
                    async with session.post(
                        self.config_loader.server_url,
                        headers=headers,
                        json=request_body,
                        timeout=self.timeout
                    ) as response:
                        # 打印响应头信息
                        # print(f"\n响应头信息:")
                        # for key, value in response.headers.items():
                        #     print(f"  {key}: {value}")
                        
                        if response.status != 200:
                            error_text = await response.text()
                            print(f"[{datetime.now().strftime('%H:%M:%S')}] 请求 {request_id} 失败:")
                            print(f"状态码: {response.status}")
                            print(f"错误信息: {error_text}")
                            return None
                        
                        response_content = ""
                        if self.config_loader.stream_enabled:
                            print_buffer = []
                            log_prefix = f"[{datetime.now().strftime('%H:%M:%S')}] 请求 {request_id}@{self.config_loader.server_url}"
                            
                            try:
                                async for line in response.content:
                                    line = line.decode('utf-8').strip()
                                    if not line:
                                        continue
                                    
                                    if line.startswith('data: '):
                                        if line == 'data: [DONE]':
                                            break
                                        try:
                                            chunk = json.loads(line[6:])
                                            choice = chunk.get('choices', [{}])[0]
                                            content = choice.get('delta', {}).get('content') or choice.get('message', {}).get('content')
                                            
                                            if content:
                                                received_tokens += len(content)
                                                response_content += content
                                                print_buffer.append(f"{log_prefix}: {content}")
                                                
                                        except (json.JSONDecodeError, KeyError, IndexError) as e:
                                            error_type = type(e).__name__
                                            print(f"{log_prefix} 解析失败 ({error_type}): {str(e)}")
                                            print(f"原始数据: {line[6:]}")
                            except asyncio.TimeoutError:
                                print(f"[{datetime.now().strftime('%H:%M:%S')}] 请求 {request_id} 流式读取超时")
                                return None
                        else:
                            try:
                                response_json = await response.json()
                                choice = response_json.get('choices', [{}])[0]
                                message = choice.get('message', {})
                                content = message.get('content', '')
                                
                                if content:
                                    response_content = content
                                    received_tokens = len(content)
                                    print(f"[{datetime.now().strftime('%H:%M:%S')}] 请求 {request_id} 完整响应接收")
                            except (json.JSONDecodeError, KeyError) as e:
                                print(f"非流式响应解析失败: {str(e)}")
                                # print(f"原始响应: {await response.text()}")
                            except asyncio.TimeoutError:
                                print(f"[{datetime.now().strftime('%H:%M:%S')}] 请求 {request_id} 响应读取超时")
                                return None
                        
                        print()
                except aiohttp.ClientConnectorError as e:
                    print(f"[{datetime.now().strftime('%H:%M:%S')}] 请求 {request_id} 连接错误:")
                    print(f"错误类型: {type(e).__name__}")
                    print(f"错误信息: {str(e)}")
                    print(f"请检查服务器地址 {self.config_loader.server_url} 是否正确，以及服务器是否可访问")
                    return None
                except aiohttp.ServerTimeoutError as e:
                    print(f"[{datetime.now().strftime('%H:%M:%S')}] 请求 {request_id} 服务器超时:")
                    print(f"错误类型: {type(e).__name__}")
                    print(f"错误信息: {str(e)}")
                    return None
                    
        except asyncio.TimeoutError:
            print(f"[{datetime.now().strftime('%H:%M:%S')}] 请求 {request_id} 连接超时")
            return None
        except aiohttp.ClientError as e:
            print(f"[{datetime.now().strftime('%H:%M:%S')}] 请求 {request_id} 发生错误:")
            print(f"错误类型: {type(e).__name__}")
            print(f"错误信息: {str(e)}")
            return None
        except Exception as e:
            print(f"[{datetime.now().strftime('%H:%M:%S')}] 请求 {request_id} 发生未知错误:")
            print(f"错误类型: {type(e).__name__}")
            print(f"错误信息: {str(e)}")
            return None
        
        elapsed_time = time.time() - start_time
        
        result = {
            "request_id": request_id,
            "elapsed_time": elapsed_time,
            "token_count": received_tokens,
            "tokens_per_second": received_tokens / elapsed_time if elapsed_time > 0 else 0,
            "response_text": response_content,
            "request_params": request_body
        }
        
        if self.config_loader.stream_enabled:
            if print_buffer:
                print(f"[{datetime.now().strftime('%H:%M:%S')}] 请求 {request_id} 流式过程接收 {len(print_buffer)} 个数据块")
            else:
                print(f"[{datetime.now().strftime('%H:%M:%S')}] 请求 {request_id} 未接收到有效数据块")
        
        return result

    async def run_benchmark(self, concurrency, questions):
        """Run benchmark with specified concurrency level"""
        print(f"\n{'=' * 50}")
        print(f"开始运行 {concurrency} 并发请求的基准测试")
        print(f"{'=' * 50}\n")
        
        semaphore = asyncio.Semaphore(concurrency)
        
        async with aiohttp.ClientSession(timeout=self.timeout) as session:
            tasks = []
            for i in range(concurrency):
                task = self.send_request(session, questions[i], i+1, semaphore)
                tasks.append(task)
            
            start_time = time.time()
            results = await asyncio.gather(*tasks, return_exceptions=True)
            total_time = time.time() - start_time
            
            # 过滤掉异常和None结果
            results = [r for r in results if r is not None and not isinstance(r, Exception)]
            
            if not results:
                print(f"并发数 {concurrency} 的所有请求都失败了")
                return None
            
            total_tokens = sum(r["token_count"] for r in results)
            avg_tokens_per_request = total_tokens / len(results)
            tokens_per_second = total_tokens / total_time if total_time > 0 else 0
            tokens_per_thread = tokens_per_second / concurrency if concurrency > 0 else 0
            
            print(f"\n{'=' * 50}")
            print(f"{concurrency} 并发请求的基准测试结果:")
            print(f"{'=' * 50}")
            print(f"总耗时: {total_time:.2f}秒")
            print(f"总生成token数: {total_tokens}")
            print(f"平均每请求token数: {avg_tokens_per_request:.2f}")
            print(f"吞吐量: {tokens_per_second:.2f} tokens/秒")
            print(f"单请求处理效率: {tokens_per_thread:.2f} tokens/秒")
            print(f"{'=' * 50}\n")
            
            return {
                "concurrency": concurrency,
                "total_time": total_time,
                "total_tokens": total_tokens,
                "avg_tokens_per_request": avg_tokens_per_request,
                "tokens_per_second": tokens_per_second,
                "tokens_per_thread": tokens_per_thread,
                "responses": [r['response_text'] for r in results],
                "requests": [r['request_params'] for r in results]
            } 